{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from PIL import Image\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import tqdm\n",
    "from torchvision.utils import save_image\n",
    "import itertools\n",
    "from torchvision import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImagePairDataset(Dataset):\n",
    "    def __init__(self, folder1, folder2, pairs, transform=None, enhance_transform_1 = None, \n",
    "                 enhance_transform_2 = None, enhance_transform_3 = None):\n",
    "        self.folder1 = folder1\n",
    "        self.folder2 = folder2\n",
    "        self.pairs = pairs\n",
    "        self.transform = transform\n",
    "        self.enhance_transform_1 = enhance_transform_1\n",
    "        self.enhance_transform_2 = enhance_transform_2\n",
    "        self.enhance_transform_3 = enhance_transform_3\n",
    "        self.image_pairs = self.read_image_pairs()\n",
    "\n",
    "    def read_image_pairs(self):\n",
    "      image_pairs = []\n",
    "      for image_pair in tqdm.tqdm(self.pairs):\n",
    "        img1_path = os.path.join(self.folder1, image_pair[0])\n",
    "        img2_path = os.path.join(self.folder2, image_pair[1])\n",
    "        img1 = Image.open(img1_path).convert(\"RGB\")\n",
    "        img2 = Image.open(img2_path).convert(\"RGB\")\n",
    "        if self.transform:\n",
    "            img9 = self.transform(img1)\n",
    "            img10 = self.transform(img2)\n",
    "            image_pairs.append((img9, img10))\n",
    "        if self.enhance_transform_1:\n",
    "            img3 = self.enhance_transform_1(img1)\n",
    "            img4 = self.enhance_transform_1(img2)\n",
    "            image_pairs.append((img3, img4))\n",
    "        if self.enhance_transform_2:\n",
    "            img5 = self.enhance_transform_2(img1)\n",
    "            img6 = self.enhance_transform_2(img2)\n",
    "            image_pairs.append((img5, img6))\n",
    "        if self.enhance_transform_3:\n",
    "            img7 = self.enhance_transform_3(img1)\n",
    "            img8 = self.enhance_transform_3(img2)\n",
    "            image_pairs.append((img7, img8))\n",
    "            \n",
    "      return image_pairs\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.pairs)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "      return self.image_pairs[idx][0], self.image_pairs[idx][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Resize((128, 256)),\n",
    "    # transforms.Grayscale(num_output_channels=3),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))  # 归一化\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "digital_dir = \"/root/autodl-tmp/Kodak_gold_200/Kodak_gold_200/digital\"\n",
    "film_dir = \"/root/autodl-tmp/Kodak_gold_200/Kodak_gold_200/film\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_pairs(image1_dir, imgage2_dir):\n",
    "    # Get the list of files in both folders\n",
    "    images1 = sorted(os.listdir(image1_dir))\n",
    "    images2 = sorted(os.listdir(imgage2_dir))\n",
    "    \n",
    "    print(len(images1))\n",
    "    print(len(images2))\n",
    "\n",
    "    # Ensure the number of files match\n",
    "    if len(images1) != len(images2):\n",
    "        raise ValueError(\"The two folders must have the same number of images.\")\n",
    "\n",
    "    # Create pairs of images (file1, file2)\n",
    "    pairs = list(zip(images1, images2))\n",
    "    \n",
    "    return pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1002\n",
      "1002\n"
     ]
    }
   ],
   "source": [
    "pairs = make_pairs(digital_dir, film_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1002/1002 [00:39<00:00, 25.36it/s]\n"
     ]
    }
   ],
   "source": [
    "# Create datasets\n",
    "train_dataset = ImagePairDataset(digital_dir, film_dir, pairs, transform = transform)\n",
    "\n",
    "# Create dataloaders\n",
    "train_loader = DataLoader(train_dataset, batch_size = batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNetGenerator(nn.Module):\n",
    "    def __init__(self, input_nc=3, output_nc=3, pretrained=True):\n",
    "        super(ResNetGenerator, self).__init__()\n",
    "\n",
    "        # Loading pre-trained ResNet-18\n",
    "        resnet = models.resnet18()\n",
    "        resnet.load_state_dict(torch.load(\"./resnet18-5c106cde.pth\"))\n",
    "        \n",
    "        # Encoder: Get the output of each layer of ResNet and construct U-Net style skip connections\n",
    "        self.input_conv = nn.Sequential(\n",
    "            resnet.conv1,   # Output: 64 channels, H/2 x W/2\n",
    "            resnet.bn1,\n",
    "            resnet.relu\n",
    "        )\n",
    "        self.maxpool = resnet.maxpool   # H/4 x W/4\n",
    "        self.encoder1 = resnet.layer1   # Output: 64 channels, H/4 x W/4\n",
    "        self.encoder2 = resnet.layer2   # Output: 128 channels, H/8 x W/8\n",
    "        self.encoder3 = resnet.layer3   # Output: 256 channels, H/16 x W/16\n",
    "        self.encoder4 = resnet.layer4   # Output: 512 channels, H/32 x W/32\n",
    "\n",
    "        # Decoder: Use upsampling (bilinear interpolation) + convolution to restore details and fuse with skip connection features\n",
    "        self.up1 = self._up_block(512, 256)  # H/32 -> H/16\n",
    "        self.up2 = self._up_block(512, 128)  # 256(from up1)+256(encoder3) -> H/16 -> H/8\n",
    "        self.up3 = self._up_block(256, 64)   # 128(from up2)+128(encoder2) -> H/8 -> H/4\n",
    "        self.up4 = self._up_block(128, 64)   # 64(from up3)+64(encoder1) -> H/4 -> H/2\n",
    "        self.up5 = self._up_block(128, 64)   # 64(from up4)+64(from input_conv) -> H/2 -> H\n",
    "\n",
    "        self.final_conv = nn.Conv2d(64, output_nc, kernel_size=1)\n",
    "\n",
    "        # Added perceptual loss module (VGG16 feature extraction)\n",
    "        self.vgg = models.vgg16()\n",
    "        self.vgg.load_state_dict(torch.load(\"./vgg16-397923af.pth\"))\n",
    "        self.vgg  = self.vgg.features[:16].eval()\n",
    "        for param in self.vgg.parameters():\n",
    "            param.requires_grad = False  # Freeze VGG weights\n",
    "\n",
    "    def _up_block(self, in_channels, out_channels):\n",
    "        \"\"\"\n",
    "        Upsampling block: upsampling first (bilinear), then convolution + ReLU\n",
    "        \"\"\"\n",
    "        block = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        return block\n",
    "\n",
    "    def add_film_grain(self, x, noise_strength=0.1):\n",
    "        \"\"\"\n",
    "        Add Gaussian noise to the image Tensor to simulate film grain effect.\n",
    "        \"\"\"\n",
    "        batch_size, channels, height, width = x.shape\n",
    "        mean = 0\n",
    "        sigma = noise_strength * 255  # Control the standard deviation of noise\n",
    "        \n",
    "        # Generate noise and overlay it on the image\n",
    "        noise = torch.normal(mean, sigma, size=(batch_size, channels, height, width), dtype=torch.float32, device=x.device)\n",
    "        noisy_image = x + noise  # Adding Noise\n",
    "        return noisy_image\n",
    "\n",
    "    def perceptual_loss(self, real, fake):\n",
    "        \"\"\"\n",
    "        Compute perceptual loss, based on VGG feature extraction.\n",
    "        \"\"\"\n",
    "        real_features = self.vgg(real)\n",
    "        fake_features = self.vgg(fake)\n",
    "        return torch.nn.functional.mse_loss(real_features, fake_features)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = x.shape[2]\n",
    "        w = x.shape[3]\n",
    "\n",
    "        # Encoder part\n",
    "        x0 = self.input_conv(x)       # x0: [B, 64, H/2, W/2]\n",
    "        x1 = self.maxpool(x0)         # x1: [B, 64, H/4, W/4]\n",
    "        x1 = self.encoder1(x1)        # x1: [B, 64, H/4, W/4]\n",
    "        x2 = self.encoder2(x1)        # x2: [B, 128, H/8, W/8]\n",
    "        x3 = self.encoder3(x2)        # x3: [B, 256, H/16, W/16]\n",
    "        x4 = self.encoder4(x3)        # x4: [B, 512, H/32, W/32]\n",
    "\n",
    "        # Decoder part: gradually upsample and concatenate with the corresponding layer features of the encoder\n",
    "        d1 = self.up1(x4)             # d1: [B, 256, H/16, W/16]\n",
    "        d1 = torch.cat([d1, x3], dim=1)  # [B, 256+256=512, H/16, W/16]\n",
    "\n",
    "        d2 = self.up2(d1)             # d2: [B, 128, H/8, W/8]\n",
    "        d2 = torch.cat([d2, x2], dim=1)  # [B, 128+128=256, H/8, W/8]\n",
    "\n",
    "        d3 = self.up3(d2)             # d3: [B, 64, H/4, W/4]\n",
    "        d3 = torch.cat([d3, x1], dim=1)  # [B, 64+64=128, H/4, W/4]\n",
    "\n",
    "        d4 = self.up4(d3)             # d4: [B, 64, H/2, W/2]\n",
    "        d4 = torch.cat([d4, x0], dim=1)  # [B, 64+64=128, H/2, W/2]\n",
    "\n",
    "        d5 = self.up5(d4)             # d5: [B, 64, H, W]\n",
    "\n",
    "        out = self.final_conv(d5)     # out: [B, output_nc, H, W]\n",
    "        out = torch.tanh(out)        \n",
    "        \n",
    "        out = self.add_film_grain(out, noise_strength=0.0002)\n",
    "\n",
    "        return out, self.perceptual_loss(x, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================\n",
    "# 3. Defining the Discriminator (PatchGAN)\n",
    "# ============================\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_nc):\n",
    "        \"\"\"\n",
    "        Using 70×70 PatchGAN discriminator\n",
    "        \"\"\"\n",
    "        super(Discriminator, self).__init__()\n",
    "        model = [\n",
    "            nn.Conv2d(input_nc, 64, kernel_size=4, stride=2, padding=1),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        ]\n",
    "\n",
    "        in_channels = 64\n",
    "        out_channels = in_channels * 2\n",
    "        # Add several layers of convolution\n",
    "        for _ in range(3):\n",
    "            model += [\n",
    "                nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1),\n",
    "                nn.InstanceNorm2d(out_channels),\n",
    "                nn.LeakyReLU(0.2, inplace=True)\n",
    "            ]\n",
    "            in_channels = out_channels\n",
    "            out_channels = in_channels * 2\n",
    "\n",
    "        # The last convolution layer\n",
    "        model += [\n",
    "            nn.Conv2d(in_channels, 1, kernel_size=4, padding=1)\n",
    "        ]\n",
    "        self.model = nn.Sequential(*model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# ============================\n",
    "# 4. Define the CycleGAN model class (add AMP support)\n",
    "# ============================\n",
    "class CycleGAN:\n",
    "    def __init__(self, device, G_AB = \"\", G_BA = \"\", D_A = \"\", D_B = \"\"):\n",
    "        self.device = device\n",
    "\n",
    "        # Initialize two generators: G_AB (digital → film), G_BA (film → digital)\n",
    "        self.G_AB = ResNetGenerator(3, 3).to(device)\n",
    "        if G_AB != \"\":\n",
    "            self.G_AB.load_state_dict(torch.load(G_AB))\n",
    "        self.G_BA = ResNetGenerator(3, 3).to(device)\n",
    "        if G_BA != \"\":\n",
    "            self.G_BA.load_state_dict(torch.load(G_BA))\n",
    "        # Initialize two discriminators: D_A (discriminates real digital images) and D_B (discriminates real film images)\n",
    "        self.D_A = Discriminator(3).to(device)\n",
    "        if D_A != \"\":\n",
    "            self.D_A.load_state_dict(torch.load(D_A))\n",
    "        self.D_B = Discriminator(3).to(device)\n",
    "        if D_B != \"\":\n",
    "            self.D_B.load_state_dict(torch.load(D_B))\n",
    "        # Define loss function: adversarial loss, cycle consistency loss, identity loss\n",
    "        self.criterion_GAN = nn.MSELoss().to(device)\n",
    "        self.criterion_cycle = nn.L1Loss().to(device)\n",
    "        self.criterion_identity = nn.L1Loss().to(device)\n",
    "\n",
    "        # Optimizer (two generators share one optimizer)\n",
    "        self.optimizer_G = optim.Adam(itertools.chain(self.G_AB.parameters(), self.G_BA.parameters()),\n",
    "                                    lr=0.0002, betas=(0.5, 0.999))\n",
    "        self.optimizer_D_A = optim.Adam(self.D_A.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "        self.optimizer_D_B = optim.Adam(self.D_B.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "\n",
    "        # Creating an AMP GradScaler object\n",
    "        self.scaler_G = torch.amp.GradScaler(device=device)\n",
    "        self.scaler_D_A = torch.amp.GradScaler(device=device)\n",
    "        self.scaler_D_B = torch.amp.GradScaler(device=device)\n",
    "\n",
    "    def set_input(self, real_A, real_B):\n",
    "        self.real_A = real_A.to(self.device)\n",
    "        self.real_B = real_B.to(self.device)\n",
    "\n",
    "    def forward(self):\n",
    "        # A→B→A\n",
    "        self.fake_B, _ = self.G_AB(self.real_A)\n",
    "        self.rec_A, _ = self.G_BA(self.fake_B)\n",
    "        # B→A→B\n",
    "        self.fake_A, _ = self.G_BA(self.real_B)\n",
    "        self.rec_B, _ = self.G_AB(self.fake_A)\n",
    "\n",
    "    def backward_G(self):\n",
    "        # Identity loss: requires the generator to remain invariant on target domain images\n",
    "        self.idt_A, _ = self.G_BA(self.real_A)\n",
    "        self.loss_idt_A = self.criterion_identity(self.idt_A, self.real_A) * 5.0 + _\n",
    "        self.idt_B, _ = self.G_AB(self.real_B)\n",
    "        self.loss_idt_B = self.criterion_identity(self.idt_B, self.real_B) * 5.0 + _\n",
    "\n",
    "        # Fighting Losses\n",
    "        pred_fake_B = self.D_B(self.fake_B)\n",
    "        target_real = torch.ones_like(pred_fake_B, device=self.device)\n",
    "        loss_GAN_AB = self.criterion_GAN(pred_fake_B, target_real)\n",
    "\n",
    "        pred_fake_A = self.D_A(self.fake_A)\n",
    "        target_real = torch.ones_like(pred_fake_A, device=self.device)\n",
    "        loss_GAN_BA = self.criterion_GAN(pred_fake_A, target_real)\n",
    "\n",
    "        # Cycle consistency loss\n",
    "        loss_cycle_A = self.criterion_cycle(self.rec_A, self.real_A) * 10.0\n",
    "        loss_cycle_B = self.criterion_cycle(self.rec_B, self.real_B) * 10.0\n",
    "\n",
    "        # Total generator loss (only saved, no backward call)\n",
    "        self.loss_G = self.loss_idt_A + self.loss_idt_B + loss_GAN_AB + loss_GAN_BA + loss_cycle_A + loss_cycle_B\n",
    "\n",
    "    def backward_D_basic(self, netD, real, fake):\n",
    "        target_real = torch.ones_like(netD(real), device=self.device)\n",
    "        target_fake = torch.zeros_like(netD(fake.detach()), device=self.device)\n",
    "        loss_real = self.criterion_GAN(netD(real), target_real)\n",
    "        loss_fake = self.criterion_GAN(netD(fake.detach()), target_fake)\n",
    "        loss_D = (loss_real + loss_fake) * 0.5\n",
    "        return loss_D\n",
    "\n",
    "    def backward_D_A(self):\n",
    "        self.loss_D_A = self.backward_D_basic(self.D_A, self.real_A, self.fake_A)\n",
    "\n",
    "    def backward_D_B(self):\n",
    "        self.loss_D_B = self.backward_D_basic(self.D_B, self.real_B, self.fake_B)\n",
    "\n",
    "    def optimize_parameters(self):\n",
    "        # --------------------------\n",
    "        # Update Generator\n",
    "        # --------------------------\n",
    "        self.optimizer_G.zero_grad()\n",
    "        with torch.cuda.amp.autocast():\n",
    "            self.forward()       # Generate fake_B, fake_A, rec_A, rec_B\n",
    "            self.backward_G()    # Calculate self.loss_G\n",
    "        self.scaler_G.scale(self.loss_G).backward()\n",
    "        self.scaler_G.step(self.optimizer_G)\n",
    "        self.scaler_G.update()\n",
    "\n",
    "        # --------------------------\n",
    "        # Update the discriminator D_A\n",
    "        # --------------------------\n",
    "        self.optimizer_D_A.zero_grad()\n",
    "        with torch.cuda.amp.autocast():\n",
    "            self.backward_D_A()  # Calculate self.loss_D_A\n",
    "        self.scaler_D_A.scale(self.loss_D_A).backward()\n",
    "        self.scaler_D_A.step(self.optimizer_D_A)\n",
    "        self.scaler_D_A.update()\n",
    "\n",
    "        # --------------------------\n",
    "        # Update the discriminator D_B\n",
    "        # --------------------------\n",
    "        self.optimizer_D_B.zero_grad()\n",
    "        with torch.cuda.amp.autocast():\n",
    "            self.backward_D_B()  # Calculate self.loss_D_B\n",
    "        self.scaler_D_B.scale(self.loss_D_B).backward()\n",
    "        self.scaler_D_B.step(self.optimizer_D_B)\n",
    "        self.scaler_D_B.update()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================\n",
    "# 5. Training loop function\n",
    "# ============================\n",
    "def train(cyclegan, dataloader, num_epochs=200, save_interval=10, pre_epoch = 0):\n",
    "    for epoch in tqdm.tqdm(range(num_epochs)):\n",
    "        for i, (data_A, data_B) in enumerate(dataloader):\n",
    "            real_A = data_A\n",
    "            real_B = data_B\n",
    "            cyclegan.set_input(real_A, real_B)\n",
    "            cyclegan.optimize_parameters()\n",
    "\n",
    "        # Save output results and models regularly\n",
    "        if (epoch + 1) % save_interval == 0:\n",
    "            os.makedirs(\"output\", exist_ok=True)\n",
    "            with torch.no_grad():\n",
    "                fake_B, _ = cyclegan.G_AB(real_A.cuda())\n",
    "                fake_A, _ = cyclegan.G_BA(real_B.cuda())\n",
    "            save_image(fake_B, f\"output/Kodak_gold_200_3/fake_B_epoch_{epoch+1+pre_epoch}.png\", normalize=True)\n",
    "            save_image(fake_A, f\"output/Kodak_gold_200_3/fake_A_epoch_{epoch+1+pre_epoch}.png\", normalize=True)\n",
    "            torch.save(cyclegan.G_AB.state_dict(), f\"output/Kodak_gold_200_3/G_AB_epoch_{epoch+1+pre_epoch}.pth\")\n",
    "            torch.save(cyclegan.G_BA.state_dict(), f\"output/Kodak_gold_200_3/G_BA_epoch_{epoch+1+pre_epoch}.pth\")\n",
    "            torch.save(cyclegan.D_A.state_dict(), f\"output/Kodak_gold_200_3/D_A_epoch_{epoch+1+pre_epoch}.pth\")\n",
    "            torch.save(cyclegan.D_B.state_dict(), f\"output/Kodak_gold_200_3/D_B_epoch_{epoch+1+pre_epoch}.pth\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1193/2430728245.py:7: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  resnet.load_state_dict(torch.load(\"./resnet18-5c106cde.pth\"))\n",
      "/tmp/ipykernel_1193/2430728245.py:32: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  self.vgg.load_state_dict(torch.load(\"./vgg16-397923af.pth\"))\n"
     ]
    }
   ],
   "source": [
    "cyclegan = CycleGAN(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1186/2430728245.py:7: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  resnet.load_state_dict(torch.load(\"./resnet18-5c106cde.pth\"))\n",
      "/tmp/ipykernel_1186/2430728245.py:32: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  self.vgg.load_state_dict(torch.load(\"./vgg16-397923af.pth\"))\n",
      "/tmp/ipykernel_1186/2374480379.py:11: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  self.G_AB.load_state_dict(torch.load(G_AB))\n",
      "/tmp/ipykernel_1186/2374480379.py:14: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  self.G_BA.load_state_dict(torch.load(G_BA))\n",
      "/tmp/ipykernel_1186/2374480379.py:18: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  self.D_A.load_state_dict(torch.load(D_A))\n",
      "/tmp/ipykernel_1186/2374480379.py:21: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  self.D_B.load_state_dict(torch.load(D_B))\n"
     ]
    }
   ],
   "source": [
    "cyclegan = CycleGAN(device, G_AB=\"output/Kodak_gold_200_3/G_AB_epoch_200.pth\", \n",
    "                    G_BA=\"output/Kodak_gold_200_3/G_BA_epoch_200.pth\", \n",
    "                    D_A=\"output/Kodak_gold_200_3/D_A_epoch_200.pth\", \n",
    "                    D_B=\"output/Kodak_gold_200_3/D_B_epoch_200.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1000 [00:00<?, ?it/s]/tmp/ipykernel_1193/2374480379.py:92: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.amp.autocast():\n",
      "/tmp/ipykernel_1193/2374480379.py:103: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.amp.autocast():\n",
      "/tmp/ipykernel_1193/2374480379.py:113: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  with torch.cuda.amp.autocast():\n",
      "  1%|          | 6/1000 [01:01<2:48:56, 10.20s/it]"
     ]
    }
   ],
   "source": [
    "train(cyclegan, train_loader, num_epochs=1000, pre_epoch = 0, save_interval=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.system(\"/usr/bin/shutdown\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "digitalFilm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
